/**
 * Copyright (c) 2018-2023, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: 
 *
 * Change Logs:
 * Date           Author            Notes
 * 2023-07-18     Shiroko           Init
 */

#include <nxos.h>

#ifdef CONFIG_NX_DRIVER_VIRTIO_MMIO

#undef NX_LOG_MOD_NAME
#define NX_LOG_MOD_NAME "VIRTIO-MMIO"

#include <base/log.h>

#include "virtio.h"
#include "virtio_mmio.h"
#include "virtio_mmio_private.h"

NX_PRIVATE NX_Bool alreadySetup = NX_False;

NX_PRIVATE NX_Device *
VirtioMMIOInitDevice(NX_Driver *driver, const char *dev_name, VirtioMMIODeviceInfo *mmioDeviceInfo, NX_Addr baseAddr)
{
    NX_LOG_D("Init MMIO Device 0x%x@0x%x for %s.", mmioDeviceInfo->memSize,
             mmioDeviceInfo->phyAddr, dev_name);
    NX_Device *device = NX_NULL;
    if (NX_DriverAttachDevice(driver, dev_name, &device) != NX_EOK)
    {
        NX_LOG_E("NX_DriverAttachDevice %s failed!", dev_name);
        return NX_NULL;
    }

    // Setup device extension field
    VirtioExtension *deviceExtension = NX_MemAlloc(sizeof(VirtioExtension));
    if (deviceExtension == NX_NULL) goto failed;
    device->extension = (void *) deviceExtension;
    deviceExtension->device = device;
    deviceExtension->busType = VIRTIO_BUS_TYPE_MMIO;

    VirtioMMIOExtension *mmioExtension = NX_MemAlloc(sizeof(VirtioMMIOExtension));
    if (mmioExtension == NX_NULL) goto failed;
    deviceExtension->busExt = (void *) mmioExtension;
    mmioExtension->baseAddr = baseAddr;

    NX_MemCopy(&mmioExtension->info, mmioDeviceInfo, sizeof(VirtioMMIODeviceInfo));

    NX_ListInit(&deviceExtension->list);

    return device;
failed:
    VirtioMMIOCleanUpDevice(device);
    return NX_NULL;
}

struct VirtioMMIOMappedArea
{
    NX_List list;
    NX_Addr phyBase;
    NX_Addr virBase;
    int count;
};
typedef struct VirtioMMIOMappedArea VirtioMMIOMappedArea;
NX_PRIVATE NX_LIST_HEAD(virtioMMIOMappedAreaListHead);


NX_PRIVATE NX_Addr VirtioMMIOMemoryMap(NX_Addr phyAddr, NX_Size size)
{
    NX_Addr virAddr = VirtioVirtualAddressAllocate(phyAddr);
    // Offset inside page of virtual and physical MUST be identical.
    NX_ASSERT((phyAddr & NX_PAGE_MASK) == (virAddr & NX_PAGE_MASK));
    // VirtIO-MMIO should not cross beyond one page.
    NX_ASSERT(((phyAddr + size - 1) & NX_PAGE_ADDR_MASK) == (phyAddr & NX_PAGE_ADDR_MASK));

    phyAddr = NX_PAGE_ALIGNDOWN(phyAddr);
    virAddr = NX_PAGE_ALIGNDOWN(virAddr);
    VirtioMMIOMappedArea *mappedArea;
    NX_ListForEachEntry(mappedArea, &virtioMMIOMappedAreaListHead, list)
    {
        if (mappedArea->virBase == virAddr)
        {
            if (phyAddr == mappedArea->phyBase)
            {
                mappedArea->count++;
                return virAddr;
            }
            else
                NX_PANIC("VirtIO-MMIO MMU Mapped memory conflict!");
        }
    }
    NX_MmuMapPageWithPhy(&gKernelMmu, virAddr, phyAddr, size,
                         NX_PAGE_ATTR_READ | NX_PAGE_ATTR_WRITE | NX_PAGE_ATTR_KERNEL);
    NX_MmuFlush();
    mappedArea = NX_MemAlloc(sizeof(VirtioMMIOMappedArea));
    NX_ListInit(&mappedArea->list);
    mappedArea->phyBase = phyAddr;
    mappedArea->virBase = virAddr;
    mappedArea->count = 1;
    NX_ListAddTail(&mappedArea->list, &virtioMMIOMappedAreaListHead);
    return virAddr;
}

NX_PRIVATE void VirtioMMIOMemoryUnmap(NX_Addr phyAddr, NX_Size size)
{
    NX_Addr virAddr = VirtioVirtualAddressAllocate(phyAddr);
    phyAddr = NX_PAGE_ALIGNDOWN(phyAddr);
    virAddr = NX_PAGE_ALIGNDOWN(virAddr);
    VirtioMMIOMappedArea *mappedArea;
    NX_ListForEachEntry(mappedArea, &virtioMMIOMappedAreaListHead, list)
    {
        if (mappedArea->virBase == virAddr)
        {
            mappedArea->count--;
            if (mappedArea->count == 0)
            {
                NX_MmuUnmapPage(&gKernelMmu, virAddr, size, NX_True);
                NX_MmuFlush();
                NX_ListDel(&mappedArea->list);
                NX_MemFree(mappedArea);
            }
            return;
        }
    }
    NX_PANIC("VirtIO-MMIO Trying to unmap a non-existed page!");
}

/* MMIO Setup */

NX_PRIVATE NX_Error VirtioMMIOSetupISR(NX_Device *device);
NX_PRIVATE void VirtioMMIORemoveISR(NX_Device *device);

struct VirtioMMIODeviceInfoList
{
    NX_List list;
    VirtioMMIODeviceInfo info;
};
typedef struct VirtioMMIODeviceInfoList VirtioMMIODeviceInfoList;
NX_PRIVATE NX_LIST_HEAD(virtioMMIODeviceInfoListHead); // Use and dispose

void VirtioMMIOCleanUpDevice(NX_Device *device)
{
    char deviceName[NX_DEVICE_NAME_LEN];
    NX_MemCopy(deviceName, device->name, NX_DEVICE_NAME_LEN);

    NX_Driver *driver = device->driver;
    if (driver == NX_NULL) return;

    VirtioExtension *deviceExtension = VirtioGetExtension(device);
    if (deviceExtension == NX_NULL) goto final;

    VirtioMMIOExtension *mmioExtension = VirtioGetBusExtension(device, VirtioMMIOExtension);
    if (mmioExtension == NX_NULL)
    {
        NX_MemFree(deviceExtension);
        goto final;
    }

    NX_Addr baseAddr = mmioExtension->info.phyAddr;
    NX_Size memSize = mmioExtension->info.memSize;
    VirtioMMIOMemoryUnmap(baseAddr, memSize);
    VirtioMMIORemoveISR(device);

    NX_ASSERT(deviceExtension->queue == NX_NULL); // Must free queue before do clean up

    NX_MemFree(mmioExtension);
    NX_MemFree(deviceExtension);
final:
    NX_DriverDetachDevice(driver, device->name);
    NX_LOG_I("Device %s detached.", deviceName);
}

NX_Error VirtioMMIOSetup(VirtioDriverTableEntity **driversTable)
{
    alreadySetup = NX_True;
    VirtioMMIODeviceInfoList *mmioDevice, *next;
    NX_ListForEachEntrySafe(mmioDevice, next, &virtioMMIODeviceInfoListHead, list)
    {
        NX_Addr phyAddr = mmioDevice->info.phyAddr;
        NX_Size memSize = mmioDevice->info.memSize;
        NX_Addr baseAddr = VirtioMMIOMemoryMap(phyAddr, memSize);
        if (VIRTIO_MMIO_REG_READ(MAGIC_VALUE, baseAddr) != VIRTIO_MMIO_MAGIC) continue;
        if (VIRTIO_MMIO_REG_READ(VERSION, baseAddr) != 2) continue; // Only modern devices are supported yet
        NX_U32 deviceId = VIRTIO_MMIO_REG_READ(DEVICE_ID, baseAddr);
        if (deviceId == 0) continue; // This is not an error according to specification, but we must ignore it
        if (deviceId >= VIRTIO_DEVICE_TYPE_SIZE)
        {
            NX_LOG_E("Device ID %d is illegal.", deviceId);
            continue;
        }
        VirtioDriverTableEntity *entity = driversTable[deviceId];
        if (entity)
        {
            // We have the corresponding driver registered.
            char deviceName[NX_DEVICE_NAME_LEN];
            NX_SNPrintf(deviceName, NX_DEVICE_NAME_LEN, entity->deviceNameTemplate, entity->deviceCount);
            NX_LOG_D("Creating device %s...", deviceName);
            NX_Device *device = VirtioMMIOInitDevice(entity->driver, deviceName, &mmioDevice->info, baseAddr);
            if (device)
            {
                // Create device successful
                entity->deviceCount++;
                NX_ListAddTail(&entity->deviceListHead,
                               &(VirtioGetExtension(device)->list));
                if (VirtioMMIOSetupISR(device) != NX_EOK)
                    NX_LOG_W("ISR Setup for device %s failed. Features may not be functional.", deviceName);
                if (entity->initHandler)
                    if (entity->initHandler(entity->driver, device) != NX_EOK)
                    {
                        NX_LOG_E("Initialize %s failed.", deviceName);
                        VirtioMMIOCleanUpDevice(device);
                        return NX_EFAULT;
                    }
                NX_LOG_I("Create %s successfully.", deviceName);
            }
        }
        else
        {
            NX_LOG_W("Driver for DeviceID %d is missing!", deviceId);
        }
        // Free node
        NX_ListDelInit(&mmioDevice->list);
        NX_MemFree(mmioDevice);
    }
    return NX_EOK;
}

// This is normally be called by DeviceTree handler or Boot command line handler
NX_Error VirtioMMIORegisterDevice(VirtioMMIODeviceInfo *info)
{
    /*
     * Device Registration MUST before Device setup.
     * MMIO Provide no hot-plug, every device is fixed on boot time.
     *
     * DeviceInfo MUST be a valid virtio-mmio device.
     */
    NX_ASSERT(alreadySetup == NX_False);

    NX_ASSERT(info->memSize != 0 && info->phyAddr != 0);
    VirtioMMIODeviceInfoList *mmioDeviceInfoList = NX_MemAlloc(sizeof(VirtioMMIODeviceInfoList));
    NX_ASSERT(mmioDeviceInfoList != NX_NULL);
    NX_MemCopy(&mmioDeviceInfoList->info, info, sizeof(VirtioMMIODeviceInfo));
    NX_ListInit(&mmioDeviceInfoList->list);
    NX_ListAddTail(&mmioDeviceInfoList->list, &virtioMMIODeviceInfoListHead);
    return NX_EOK;
}

/* MMIO Common Cfg */
void VirtioMMIOCommonConfigGet(NX_Device *device, VirtioCommonConfigItem item, void *value)
{
    NX_Addr base = VirtioGetBusExtension(device, VirtioMMIOExtension)->baseAddr;
    switch (item)
    {
    case VIRTIO_COMMON_CFG_FEATURE_DEVICE:
        VIRTIO_MMIO_REG_WRITE(DEVICE_FEATURES_SEL, base, ((VirtioFeature *) value)->selector);
        NX_MemoryBarrier();
        ((VirtioFeature *) value)->features = VIRTIO_MMIO_REG_READ(DEVICE_FEATURES, base);
        break;
    case VIRTIO_COMMON_CFG_FEATURE_DRIVER:
        NX_LOG_E("Feature Driver of Common Config is Write-only.");
        *((NX_U32 *) value) = 0;
        break;
    case VIRTIO_COMMON_CFG_DEVICE_STATUS:
        *((NX_U8 *) value) = (NX_U8) (VIRTIO_MMIO_REG_READ(STATUS, base) & 0xFF);
        break;
    case VIRTIO_COMMON_CFG_QUEUE_NUM_MAX:
        // num_queues field of PCI have no corresponding field in MMIO
        *((NX_U16 *) value) = (NX_U16) (0xFFFF);
        break;
    case VIRTIO_COMMON_CFG_QUEUE_SELECT:
        NX_LOG_E("Queue Select of Common Config is Write-only.");
        *((NX_U16 *) value) = 0;
        break;
    case VIRTIO_COMMON_CFG_QUEUE_DESC:
    case VIRTIO_COMMON_CFG_QUEUE_DRIVER:
    case VIRTIO_COMMON_CFG_QUEUE_DEVICE:
        NX_LOG_E("Queue_Desc/Driver/Device of Common Config is Write-only.");
        *((NX_U64 *) value) = 0;
        break;
    case VIRTIO_COMMON_CFG_QUEUE_ENABLE:
        *((NX_U16 *) value) = (NX_U16) (VIRTIO_MMIO_REG_READ(QUEUE_READY, base) & 0xFFFF);
        break;
    case VIRTIO_COMMON_CFG_QUEUE_SIZE:
        *((NX_U16 *) value) = (NX_U16) (VIRTIO_MMIO_REG_READ(QUEUE_NUM_MAX, base) & 0xFFFF);
        break;
    case VIRTIO_COMMON_CFG_CONFIG_GENERATION:
        *((NX_U8 *) value) = (NX_U16) (VIRTIO_MMIO_REG_READ(CONFIG_GENERATION, base) & 0xFF);
        break;
    }
}

void VirtioMMIOCommonConfigSet(NX_Device *device, VirtioCommonConfigItem item, void *value)
{
    NX_Addr base = VirtioGetBusExtension(device, VirtioMMIOExtension)->baseAddr;
    switch (item)
    {
    case VIRTIO_COMMON_CFG_FEATURE_DEVICE:
        NX_PANIC("Feature_Device of Common Config is Read-only.");
        break;
    case VIRTIO_COMMON_CFG_FEATURE_DRIVER:
        VIRTIO_MMIO_REG_WRITE(DRIVER_FEATURES_SEL, base, ((VirtioFeature *) value)->selector);
        NX_MemoryBarrier();
        VIRTIO_MMIO_REG_WRITE(DRIVER_FEATURES, base, ((VirtioFeature *) value)->features);
        break;
    case VIRTIO_COMMON_CFG_DEVICE_STATUS:
        VIRTIO_MMIO_REG_WRITE(STATUS, base, (NX_U32) (*((NX_U8 *) value)));
        break;
    case VIRTIO_COMMON_CFG_QUEUE_NUM_MAX:
        NX_PANIC("Queue_Num_Max of Common Config is Read-only.");
        break;
    case VIRTIO_COMMON_CFG_QUEUE_SELECT:
        VIRTIO_MMIO_REG_WRITE(QUEUE_SEL, base, (NX_U32) (*((NX_U16 *) value)));
        break;
    case VIRTIO_COMMON_CFG_QUEUE_DESC:
    case VIRTIO_COMMON_CFG_QUEUE_DRIVER:
    case VIRTIO_COMMON_CFG_QUEUE_DEVICE:
    {
        NX_U64 addr = *((NX_U64 *) value);
        NX_U32 low = (NX_U32) (addr & 0xFFFFFFFF);
        NX_U32 high = (NX_U32) (addr >> 32);
        switch (item)
        {
        case VIRTIO_COMMON_CFG_QUEUE_DESC:
            VIRTIO_MMIO_REG_WRITE(QUEUE_DESC_LOW, base, low);
            VIRTIO_MMIO_REG_WRITE(QUEUE_DESC_HIGH, base, high);
            break;
        case VIRTIO_COMMON_CFG_QUEUE_DRIVER:
            VIRTIO_MMIO_REG_WRITE(QUEUE_DRIVER_LOW, base, low);
            VIRTIO_MMIO_REG_WRITE(QUEUE_DRIVER_HIGH, base, high);
            break;
        case VIRTIO_COMMON_CFG_QUEUE_DEVICE:
            VIRTIO_MMIO_REG_WRITE(QUEUE_DEVICE_LOW, base, low);
            VIRTIO_MMIO_REG_WRITE(QUEUE_DEVICE_HIGH, base, high);
            break;
        default:
            NX_PANIC("Impossible branch of switch");
            break;
        }
        break;
    }
    case VIRTIO_COMMON_CFG_QUEUE_ENABLE:
        VIRTIO_MMIO_REG_WRITE(QUEUE_READY, base, (*((NX_U16 *) value) & 0xFFFF));
        break;
    case VIRTIO_COMMON_CFG_QUEUE_SIZE:
        VIRTIO_MMIO_REG_WRITE(QUEUE_NUM, base, (*((NX_U16 *) value) & 0xFFFF));
        break;
    case VIRTIO_COMMON_CFG_CONFIG_GENERATION:
        NX_PANIC("Virtio Config Generation is Read-Only.");
        break;
    }
}

/* MMIO Device Cfg */
void *VirtioMMIOGetDeviceConfig(NX_Device *device)
{
    // This is actually risky with no volatile guard on memory io.
    return (void *) (VirtioGetBusExtension(device, VirtioMMIOExtension)->baseAddr +
                     VIRTIO_MMIO_REG_OFFSET_DEVICE_CONFIG);
}

/* MMIO Send Available Notification */
void VirtioMMIOSendNotification(NX_Device *device, NX_U16 queue_index)
{
    NX_Addr base = VirtioGetBusExtension(device, VirtioMMIOExtension)->baseAddr;
    // VIRTIO_F_NOTIFICATION_DATA should not be negotiated.
    VIRTIO_MMIO_REG_WRITE(QUEUE_NOTIFY, base, queue_index);
}

/* Interrupt Service Routine */

NX_PRIVATE NX_Error VirtioMMIO_ISR(NX_IRQ_Number irqNumber, void *data)
{
    NX_Device *device = (NX_Device *) data;
    VirtioMMIOExtension *mmioExtension = VirtioGetBusExtension(device, VirtioMMIOExtension);
    NX_U32 isrStatus = VIRTIO_MMIO_REG_READ(INTERRUPT_STATUS, mmioExtension->baseAddr);
    NX_Error err = NX_EINVAL;
    if (isrStatus & 1)
    {
        // Bit 0 set, queue interrupt
        err = VirtioISR(VIRTIO_ISR_TYPE_QUEUE, device);
        VIRTIO_MMIO_REG_WRITE(INTERRUPT_ACK, mmioExtension->baseAddr, 1);
    }
    else if (isrStatus & 2)
    {
        // Bit 1 set, configuration interrupt
        err = VirtioISR(VIRTIO_ISR_TYPE_CONFIG, device);
        VIRTIO_MMIO_REG_WRITE(INTERRUPT_ACK, mmioExtension->baseAddr, 2);
    }
    return err;
}

NX_PRIVATE NX_Error VirtioMMIOSetupISR(NX_Device *device)
{
    VirtioSetupISR(device);
    VirtioMMIOExtension *mmioExtension = VirtioGetBusExtension(device, VirtioMMIOExtension);
    if (NX_IRQ_Bind(mmioExtension->info.interrupt, VirtioMMIO_ISR, (void *) device, device->name, 0) != NX_EOK)
        return NX_EFAULT;
    if (NX_IRQ_Unmask(mmioExtension->info.interrupt) != NX_EOK)
        return NX_EFAULT;
    return NX_EOK;
}

NX_PRIVATE void VirtioMMIORemoveISR(NX_Device *device)
{
    VirtioMMIOExtension *mmioExtension = VirtioGetBusExtension(device, VirtioMMIOExtension);
    NX_IRQ_Unbind(mmioExtension->info.interrupt, (void *) device);
    VirtioRemoveISR(device);
}


#endif // CONFIG_NX_DRIVER_VIRTIO_MMIO
